{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Simple Arrays\n",
    "\n",
    "In this example, we will show how to start building a simple Array API.\n",
    "\n",
    "The important things to note are:\n",
    "\n",
    "1. Our \"optimization\" replacement isn't tied to any particular type of array. It is agnostic to how we actually represent arrays (in this case strings).\n",
    "2. We end up with a Python class that acts like an array, but builds up an expression graph. We use this same class to build up the replacements, so they are very succinct and readble.\n",
    "3. There is no metaprogramming in this example, and no inspecting of type signatures.\n",
    "4. Everthing is well typed and MyPy will warn us if we are trying to call operations on arrays which are not valid.\n",
    "\n",
    "We only define two operations, add and getitem:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "import metadsl\n",
    "import metadsl_rewrite\n",
    "import typing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Array(metadsl.Expression):\n",
    "    @metadsl.expression\n",
    "    def __add__(self, other: Array) -> Array:\n",
    "        ...\n",
    "    \n",
    "    @metadsl.expression\n",
    "    def __getitem__(self, idx: int) -> Array:\n",
    "        ...\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define one replacement, that distributes indexing over addition:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "register_array = metadsl_rewrite.register['array']\n",
    "\n",
    "@register_array\n",
    "@metadsl_rewrite.rule\n",
    "def _distribute_add(a: Array, b: Array, c: int):\n",
    "    return (\n",
    "        # expression to match agains\n",
    "        (a + b)[c],\n",
    "        # expression it is replaced with\n",
    "        lambda: a[c] + b[c]\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we define a simple \"implementation\" of arrays, that are just strings representing their names."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "@metadsl.expression\n",
    "def array(x: str) -> Array:\n",
    "    ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array(\"a\")"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a = array(\"a\")\n",
    "a"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We write how to convert adding and indexing into strings, to \"compile\" the expression:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "register_array_string = metadsl_rewrite.register['array']\n",
    "\n",
    "\n",
    "@metadsl.expression\n",
    "def as_string(x: Array) ->  str:\n",
    "    ...\n",
    "\n",
    "@register_array_string\n",
    "@metadsl_rewrite.rule\n",
    "def _string(x: str):\n",
    "    return as_string(array(x)), lambda: x\n",
    "\n",
    "\n",
    "@register_array_string\n",
    "@metadsl_rewrite.rule\n",
    "def _add(x: str, y: str):\n",
    "    return (\n",
    "        array(x) + array(y),\n",
    "        lambda: array(f\"({x} + {y})\")\n",
    "    )\n",
    "\n",
    "\n",
    "@register_array_string\n",
    "@metadsl_rewrite.rule\n",
    "def _getitem(x: str, y: int):\n",
    "    return (\n",
    "        array(x)[y],\n",
    "        lambda: array(f\"{x}[{y}]\")\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can create two arrays, add them and index, and turn them into a string, which will run the replacments:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(a + b)[10]\n"
     ]
    }
   ],
   "source": [
    "a = array(\"a\")\n",
    "b = array(\"b\")\n",
    "\n",
    "res = (a + b)[10]\n",
    "print(metadsl_rewrite.execute(as_string(res)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.0 ('metadsl')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.0"
  },
  "vscode": {
   "interpreter": {
    "hash": "8f1320fcc4eb9d472fe42e82e7e653b438a14df950aaaa73e390cd48905e514f"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
